import cv2
import numpy as np
import pyautogui
from ultralytics import YOLO
import time


def capture_screen(x, y, width, height):
    """
    捕获指定区域的屏幕截图
    :param x: 区域左上角 x 坐标
    :param y: 区域左上角 y 坐标
    :param width: 区域宽度
    :param height: 区域高度
    :return: 转换为 OpenCV 格式的截图
    """
    screenshot = pyautogui.screenshot(region=(x, y, width, height))
    return cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)


def detect_objects(image, model):
    """
    使用 YOLOv8 模型对图像进行目标检测
    :param image: 待检测的图像
    :param model: 加载好的 YOLO 模型
    :return: 检测结果
    """
    return model(image)


def print_detection_results(results):
    """
    打印检测结果，并返回矿石和岩浆的位置列表
    :param results: 检测结果
    :return: 矿石位置列表和岩浆位置列表
    """
    detection_info = {
        '岩浆': [],
        '石块': [],
        '自己': []
    }
    for r in results:
        boxes = r.boxes
        for box in boxes:
            class_id = r.names[int(box.cls)]
            confidence = float(box.conf)
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            print(f"检测到对象: {class_id}, 置信度: {confidence:.2f}, 位置: ({x1}, {y1}, {x2}, {y2})")
            if class_id in detection_info:
                detection_info[class_id].append((x1, y1, x2, y2))
    return detection_info['石块'], detection_info['岩浆'], detection_info['自己']


def get_center(x1, y1, x2, y2):
    """
    计算目标的中心位置，并确定以中心为基准的 10x10 像素方块的位置
    :param x1: 目标左上角 x 坐标
    :param y1: 目标左上角 y 坐标
    :param x2: 目标右下角 x 坐标
    :param y2: 目标右下角 y 坐标
    :return: 10x10 像素方块的左上角和右下角坐标
    """
    center_x = int((x1 + x2) / 2)
    center_y = int((y1 + y2) / 2)
    crop_x1 = max(center_x - 5, 0)
    crop_y1 = max(center_y - 5, 87)
    crop_x2 = min(center_x + 5, 842)  # 根据截图宽度调整
    crop_y2 = min(center_y + 5, 658)  # 根据截图高度调整
    print(f"人物方块位置: 左上角({crop_x1}, {crop_y1}) 右下角({crop_x2}, {crop_y2})")
    return crop_x1, crop_y1, crop_x2, crop_y2

def main():
    x = 0
    y = 87
    width = 842
    height = 570

    # 加载模型
    model = YOLO('best.pt')

    while True:
        # 捕获屏幕截图
        screenshot = capture_screen(x, y, width, height)
        # yolo 检测
        results = detect_objects(screenshot, model)
        # 打印检测结果并获取矿石和岩浆位置
        ore_locations, magma_locations,role_locations = print_detection_results(results)
        if ore_locations:
            print("检测到的石块位置：")
            for loc in ore_locations:
                print(f"左上角({loc[0]}, {loc[1]}) 右下角({loc[2]}, {loc[3]})")
        if magma_locations:
            print("检测到的岩浆位置：")
            for loc in magma_locations:
                print(f"左上角({loc[0]}, {loc[1]}) 右下角({loc[2]}, {loc[3]})")
        if role_locations:
            print("检测到的人物位置：")
            for loc in role_locations:
                print(f"左上角({loc[0]}, {loc[1]}) 右下角({loc[2]}, {loc[3]})")

        for r in results:
            im_array = r.plot()
            cv2.imshow('YOLOv8 Inference', im_array)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        # 可以根据实际需求调整等待时间或者移除
        # time.sleep(1)

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()

#规则：
# 1.将获得到的石块位置和岩浆位置设置为障碍物（是一个区域）
# 2.将人物边缘来作为标记区域。人物的边缘不能碰到障碍物的边缘。如果碰到了就死了。
# 3.让人物在安全区域内移动。